#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
样本数据加载器模块
负责高效加载和索引样本数据表，提供快速查询接口
"""

import os
import pandas as pd
import logging
from typing import Dict, Any, Optional, List
import threading
import hashlib
from datetime import datetime

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
)
logger = logging.getLogger(__name__)


class SampleDataLoader:
    """
    样本数据加载器
    高效加载样本数据表并提供基于视频ID的快速查询功能
    """
    
    def __init__(self, sample_data_path: str = None):
        """
        初始化样本数据加载器
        
        Args:
            sample_data_path: 样本数据文件路径
        """
        self.sample_data_path = sample_data_path
        self.sample_data = None
        self.sample_data_index = {}
        self._lock = threading.RLock()
        self._last_loaded_time = None
        self._file_hash = None
        
        # 自动加载数据（如果提供了路径）
        if sample_data_path:
            self.load_data(sample_data_path)
    
    def _calculate_file_hash(self, file_path: str) -> str:
        """
        计算文件哈希值，用于检测文件是否变更
        
        Args:
            file_path: 文件路径
            
        Returns:
            str: 文件哈希值
        """
        try:
            hasher = hashlib.md5()
            with open(file_path, 'rb') as f:
                # 分块读取文件计算哈希
                for chunk in iter(lambda: f.read(4096), b''):
                    hasher.update(chunk)
            return hasher.hexdigest()
        except Exception as e:
            logger.error(f"计算文件哈希失败: {str(e)}")
            return None
    
    def load_data(self, file_path: str = None) -> bool:
        """
        加载样本数据
        
        Args:
            file_path: 样本数据文件路径
            
        Returns:
            bool: 加载是否成功
        """
        with self._lock:
            # 如果未提供路径，使用初始化时的路径
            if file_path is None:
                file_path = self.sample_data_path
            
            if not file_path:
                logger.error("未指定样本数据文件路径")
                return False
            
            if not os.path.exists(file_path):
                logger.error(f"样本数据文件不存在: {file_path}")
                return False
            
            try:
                # 计算文件哈希，检查是否需要重新加载
                current_hash = self._calculate_file_hash(file_path)
                if current_hash and current_hash == self._file_hash:
                    logger.info(f"样本数据文件未变更，无需重新加载: {file_path}")
                    return True
                
                # 读取CSV文件
                logger.info(f"开始加载样本数据: {file_path}")
                df = pd.read_csv(file_path, encoding='utf-8')
                
                # 记录原始路径
                self.sample_data_path = file_path
                self.sample_data = df
                self._last_loaded_time = datetime.now()
                self._file_hash = current_hash
                
                # 构建索引
                self._build_index()
                
                logger.info(f"样本数据加载成功，共{len(df)}条记录")
                return True
                
            except UnicodeDecodeError:
                # 尝试其他编码
                try:
                    df = pd.read_csv(file_path, encoding='gbk')
                    self.sample_data = df
                    self._last_loaded_time = datetime.now()
                    self._build_index()
                    logger.info(f"使用GBK编码加载样本数据成功，共{len(df)}条记录")
                    return True
                except Exception as e:
                    logger.error(f"使用GBK编码加载样本数据失败: {str(e)}")
                    return False
            except Exception as e:
                logger.error(f"加载样本数据失败: {str(e)}")
                return False
    
    def _build_index(self):
        """
        构建样本数据索引，支持快速查询
        """
        if self.sample_data is None or self.sample_data.empty:
            self.sample_data_index = {}
            return
        
        # 重置索引
        self.sample_data_index = {}
        
        # 尝试识别视频ID列
        video_id_columns = ['视频ID', 'video_id', 'id', '视频编码']
        found_id_column = None
        
        # 查找有效的ID列
        for col in video_id_columns:
            if col in self.sample_data.columns:
                # 检查该列是否包含非空值
                if not self.sample_data[col].isna().all():
                    found_id_column = col
                    break
        
        if not found_id_column:
            logger.warning("未找到有效的视频ID列")
            return
        
        # 构建索引
        for idx, row in self.sample_data.iterrows():
            video_id = row[found_id_column]
            if pd.notna(video_id):
                # 确保ID是字符串类型
                video_id_str = str(video_id).strip()
                if video_id_str:
                    self.sample_data_index[video_id_str] = idx
        
        logger.info(f"构建样本数据索引完成，索引条目数: {len(self.sample_data_index)}")
    
    def get_sample_by_video_id(self, video_id: str) -> Optional[Dict[str, Any]]:
        """
        根据视频ID获取样本数据
        
        Args:
            video_id: 视频ID
            
        Returns:
            dict: 样本数据记录，如果未找到返回None
        """
        with self._lock:
            if self.sample_data is None or self.sample_data.empty:
                logger.warning("样本数据未加载")
                return None
            
            # 确保ID是字符串类型
            video_id_str = str(video_id).strip()
            
            # 首先通过索引查询
            if video_id_str in self.sample_data_index:
                idx = self.sample_data_index[video_id_str]
                return self.sample_data.iloc[idx].to_dict()
            
            # 如果索引未命中，尝试遍历查找（较慢）
            video_id_columns = ['视频ID', 'video_id', 'id', '视频编码']
            for col in video_id_columns:
                if col in self.sample_data.columns:
                    # 转换为字符串进行比较
                    matches = self.sample_data[
                        self.sample_data[col].astype(str).str.strip() == video_id_str
                    ]
                    if not matches.empty:
                        # 更新索引以加快未来查询
                        self.sample_data_index[video_id_str] = matches.index[0]
                        return matches.iloc[0].to_dict()
            
            logger.info(f"未找到视频ID {video_id} 的样本数据")
            return None
    
    def get_sample_by_index(self, index: int) -> Optional[Dict[str, Any]]:
        """
        根据索引获取样本数据
        
        Args:
            index: 样本数据索引
            
        Returns:
            dict: 样本数据记录，如果索引无效返回None
        """
        with self._lock:
            if self.sample_data is None or self.sample_data.empty:
                return None
            
            if 0 <= index < len(self.sample_data):
                return self.sample_data.iloc[index].to_dict()
            else:
                logger.warning(f"索引越界: {index}")
                return None
    
    def search_samples(self, search_term: str, columns: List[str] = None) -> List[Dict[str, Any]]:
        """
        搜索样本数据
        
        Args:
            search_term: 搜索关键词
            columns: 要搜索的列名列表，如果为None则搜索所有列
            
        Returns:
            list: 匹配的样本数据记录列表
        """
        with self._lock:
            if self.sample_data is None or self.sample_data.empty:
                return []
            
            # 确保搜索词是字符串
            search_term = str(search_term).lower()
            if not search_term:
                return []
            
            # 确定要搜索的列
            search_columns = columns
            if search_columns is None:
                # 默认只搜索文本列
                search_columns = []
                for col in self.sample_data.columns:
                    if self.sample_data[col].dtype == 'object':
                        search_columns.append(col)
            
            # 搜索匹配的行
            matches = []
            for idx, row in self.sample_data.iterrows():
                for col in search_columns:
                    if col in row and pd.notna(row[col]):
                        if search_term in str(row[col]).lower():
                            matches.append(row.to_dict())
                            break
            
            logger.info(f"搜索 '{search_term}' 找到 {len(matches)} 条匹配结果")
            return matches
    
    def get_sample_count(self) -> int:
        """
        获取样本数据总数
        
        Returns:
            int: 样本数据记录数
        """
        with self._lock:
            if self.sample_data is None:
                return 0
            return len(self.sample_data)
    
    def get_index_size(self) -> int:
        """
        获取索引大小
        
        Returns:
            int: 索引条目数
        """
        with self._lock:
            return len(self.sample_data_index)
    
    def get_last_loaded_time(self) -> Optional[datetime]:
        """
        获取最后加载时间
        
        Returns:
            datetime: 最后加载时间
        """
        with self._lock:
            return self._last_loaded_time
    
    def is_data_loaded(self) -> bool:
        """
        检查数据是否已加载
        
        Returns:
            bool: 是否已加载数据
        """
        with self._lock:
            return self.sample_data is not None and not self.sample_data.empty
    
    def refresh_data(self) -> bool:
        """
        刷新样本数据（重新加载）
        
        Returns:
            bool: 刷新是否成功
        """
        return self.load_data()
    
    def get_column_names(self) -> List[str]:
        """
        获取样本数据的列名
        
        Returns:
            list: 列名列表
        """
        with self._lock:
            if self.sample_data is None:
                return []
            return list(self.sample_data.columns)
    
    def get_sample_statistics(self) -> Dict[str, Any]:
        """
        获取样本数据统计信息
        
        Returns:
            dict: 统计信息
        """
        with self._lock:
            stats = {
                'total_records': self.get_sample_count(),
                'index_size': self.get_index_size(),
                'last_loaded_time': self._last_loaded_time.isoformat() if self._last_loaded_time else None,
                'file_path': self.sample_data_path,
                'column_count': len(self.get_column_names()),
                'columns': self.get_column_names()
            }
            
            # 添加每列的数据统计
            if self.sample_data is not None:
                column_stats = {}
                for col in self.sample_data.columns:
                    non_null_count = self.sample_data[col].count()
                    null_count = len(self.sample_data) - non_null_count
                    column_stats[col] = {
                        'non_null_count': non_null_count,
                        'null_count': null_count,
                        'non_null_percentage': round((non_null_count / len(self.sample_data)) * 100, 2) if len(self.sample_data) > 0 else 0
                    }
                stats['column_statistics'] = column_stats
            
            return stats


# 创建全局实例
sample_data_loader_instance = None

def get_sample_data_loader(sample_data_path: str = None) -> SampleDataLoader:
    """
    获取样本数据加载器实例（单例模式）
    
    Args:
        sample_data_path: 样本数据文件路径
        
    Returns:
        SampleDataLoader: 样本数据加载器实例
    """
    global sample_data_loader_instance
    
    if sample_data_loader_instance is None:
        # 如果未提供路径，使用默认路径
        if not sample_data_path:
            # 默认样本数据路径
            default_path = os.path.join(
                os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
                'data', '样本数据表.csv'
            )
            sample_data_path = default_path
        
        sample_data_loader_instance = SampleDataLoader(sample_data_path)
    
    return sample_data_loader_instance


if __name__ == '__main__':
    try:
        # 测试代码
        print("SampleDataLoader测试")
        
        # 默认路径测试
        loader = get_sample_data_loader()
        print(f"样本数据加载状态: {loader.is_data_loaded()}")
        print(f"样本数据记录数: {loader.get_sample_count()}")
        print(f"索引条目数: {loader.get_index_size()}")
        
        # 尝试获取一条样本数据
        sample = loader.get_sample_by_index(0)
        if sample:
            print("\n第一条样本数据:")
            # 只打印前几个字段
            for i, (key, value) in enumerate(sample.items()):
                if i < 5:  # 只显示前5个字段
                    print(f"  {key}: {value}")
                else:
                    print(f"  ... 还有 {len(sample) - 5} 个字段")
                    break
        
        # 打印统计信息
        print("\n样本数据统计信息:")
        stats = loader.get_sample_statistics()
        print(f"  总记录数: {stats['total_records']}")
        print(f"  列数: {stats['column_count']}")
        print(f"  最后加载时间: {stats['last_loaded_time']}")
        print(f"  文件路径: {stats['file_path']}")
        
        print("\nSampleDataLoader测试完成")
        
    except Exception as e:
        print(f"测试失败: {str(e)}")